%% This function for epoching the head-turning data

% Input:
% EEG is the preprocessed data
% evenOnsetName is the name of event onset; eventOnsetName = {'1' '2'}
% conditionsName is the list of walking condition; conditionsName = {'14' '18' '24' '28'}
% epochLength is length of epoch; lengthWalkEpoch = [0 32];
% epochDIr is the directory for the generated epoch data

function epoch(EEG, eventOnsetName, conditionsName, lengthWalkEpoch, epochDir)

disp(['EPOCHING DATASET: ' EEG.filename]);
    
eeglab_complete_path = which('eeglab');
eeglab_split_path = strsplit(eeglab_complete_path,'eeglab.m');
eeglabDir = eeglab_split_path{1};
   
%% EPOCH THE DATA INTO WALK CONDIOTIONS
EEG = pop_epoch( EEG, {  eventOnsetName  }, lengthWalkEpoch, 'epochinfo', 'yes');
EEG = pop_rmbase( EEG, [lengthWalkEpoch(1) 0]*1000);    % Change from seconds to miliseconds
EEG = eeg_checkset(EEG);

%% Segment the data for further analysis into straight1, turn and straight 2 for 60° and 85° turning trials
epochIdxS1_60 = [];
epochIdxT_60 = [];
epochIdxS2_60 = [];
epochIdxS1_85 = [];
epochIdxT_85 = [];
epochIdxS2_85 = [];

for epochCur =1:size(EEG.epoch,2)
    tempEventType = EEG.epoch(epochCur).eventtype;
    for event = 1:size(tempEventType,2)
%         first straight walking for 60° turning trial
        if (strcmp(tempEventType{event},'1')||...
                strcmp(tempEventType{event},'2'))&&...
                (strcmp(tempEventType{event},'14')||...
                strcmp(tempEventType{event},'24'))
            epochIdxS1_60 = [epochIdxS1_60; epochCur];   
            break; % CAREFUL MISSING DATA
        end
        
%         turning for 60° turning trial       
        if strcmp(tempEventType{event},'14')||...
                strcmp(tempEventType{event},'24')
            epochIdxT_60 = [epochIdxT_60; epochCur];   
            break; % CAREFUL MISSING DATA
        end
        
%         second straight walking for 60° turning trial
        if (strcmp(tempEventType{event},'4'))&&...
                (strcmp(tempEventType{event},'14')||...
                strcmp(tempEventType{event},'24'))
            epochIdxS2_60 = [epochIdxS2_60; epochCur];   
            break; % CAREFUL MISSING DATA
        end
        
%         first straight walking for 85° turning trial
        if (strcmp(tempEventType{event},'1')||...
                strcmp(tempEventType{event},'2'))&&...
                (strcmp(tempEventType{event},'18')||...
                strcmp(tempEventType{event},'28'))
            epochIdxS1_85 = [epochIdxS1_85; epochCur];   
            break; % CAREFUL MISSING DATA
        end
        
%         turning for 85° turning trial       
        if strcmp(tempEventType{event},'18')||...
                strcmp(tempEventType{event},'28')
            epochIdxT_85 = [epochIdxT_85; epochCur];   
            break; % CAREFUL MISSING DATA
        end
        
%         second straight walking for 85° turning trial
        if (strcmp(tempEventType{event},'4'))&&...
                (strcmp(tempEventType{event},'18')||...
                strcmp(tempEventType{event},'28'))
            epochIdxS2_85 = [epochIdxS2_85; epochCur];   
            break; % CAREFUL MISSING DATA
        end        
    end
end

% Set the lable to segment the data
trial_S1_60 = ones(1,size(EEG.epoch,2));
trial_S1_60(epochIdxS1_60) = 0;

trial_T_60 = ones(1,size(EEG.epoch,2));
trial_T_60(epochIdxT_60) = 0;

trial_S2_60 = ones(1,size(EEG.epoch,2));
trial_S2_60(epochIdxS2_60) = 0;

trial_S1_85 = ones(1,size(EEG.epoch,2));
trial_S1_85(epochIdxS1_85) = 0;

trial_T_85 = ones(1,size(EEG.epoch,2));
trial_T_85(epochIdxT_85) = 0;

trial_S2_85 = ones(1,size(EEG.epoch,2));
trial_S2_85(epochIdxS2_85) = 0;

%%
mkdir([epochDir '\' EEG.filename(1:3)]);

EEG_S1_60 = pop_rejepoch( EEG, trial_S1_60, 0);
EEG_S2_60 = pop_rejepoch( EEG, trial_S2_60, 0);
EEG_T_60 = pop_rejepoch( EEG, trial_T_60, 0);
EEG_S1_85 = pop_rejepoch( EEG, trial_S1_85, 0);
EEG_S2_85 = pop_rejepoch( EEG, trial_S2_85, 0);
EEG_T_85 = pop_rejepoch( EEG, trial_T_85, 0);

% automatic artifact rejection
EEG_S1_60 = pop_autorej(EEG_S1_60, 'nogui','on','threshold',1000,'eegplot','off');
EEG_T_60 = pop_autorej(EEG_T_60, 'nogui','on','threshold',1000,'eegplot','off');
EEG_S2_60 = pop_autorej(EEG_S2_60, 'nogui','on','threshold',1000,'eegplot','off');
EEG_S1_85 = pop_autorej(EEG_S1_85, 'nogui','on','threshold',1000,'eegplot','off');
EEG_T_85 = pop_autorej(EEG_T_85, 'nogui','on','threshold',1000,'eegplot','off');
EEG_S2_85 = pop_autorej(EEG_S2_85, 'nogui','on','threshold',1000,'eegplot','off');

[EEG_S1_60] = eeg_checkset(EEG_S1_60);
[EEG_S2_60] = eeg_checkset(EEG_S2_60);
[EEG_T_60] = eeg_checkset(EEG_T_60);
[EEG_S1_85] = eeg_checkset(EEG_S1_85);
[EEG_S2_85] = eeg_checkset(EEG_S2_85);
[EEG_T_85] = eeg_checkset(EEG_T_85);


%% Saving the data
[EEG_S1_60] = eeg_checkset(EEG_S1_60);
[EEG_S2_60] = eeg_checkset(EEG_S2_60);
[EEG_T_60] = eeg_checkset(EEG_T_60);
[EEG_S1_85] = eeg_checkset(EEG_S1_85);
[EEG_S2_85] = eeg_checkset(EEG_S2_85);
[EEG_T_85] = eeg_checkset(EEG_T_85);

pop_saveset( EEG_S1_60, 'filename', [EEG.filename(1:3) '_' conditionsName{1} '_Epoch'], 'filepath', [epochDir '\' EEG.filename(1:3)]);
pop_saveset( EEG_S2_60, 'filename', [EEG.filename(1:3) '_' conditionsName{2} '_Epoch'], 'filepath', [epochDir '\' EEG.filename(1:3)]);
pop_saveset( EEG_T_60, 'filename', [EEG.filename(1:3) '_' conditionsName{3} '_Epoch'], 'filepath', [epochDir '\' EEG.filename(1:3)]);
pop_saveset( EEG_S1_85, 'filename', [EEG.filename(1:3) '_' conditionsName{4} '_Epoch'], 'filepath', [epochDir '\' EEG.filename(1:3)]);
pop_saveset( EEG_S2_85, 'filename', [EEG.filename(1:3) '_' conditionsName{5} '_Epoch'], 'filepath', [epochDir '\' EEG.filename(1:3)]);
pop_saveset( EEG_T_85, 'filename', [EEG.filename(1:3) '_' conditionsName{6} '_Epoch'], 'filepath', [epochDir '\' EEG.filename(1:3)]);

disp([EEG_S1_60.filename ' has ' num2str(size(EEG_S1_60.epoch,2)) ' epochs']); 
disp([EEG_S2_60.filename ' has ' num2str(size(EEG_S2_60.epoch,2)) ' epochs']); 
disp([EEG_T_60.filename ' has ' num2str(size(EEG_T_60.epoch,2)) ' epochs']); 
disp([EEG_S1_85.filename ' has ' num2str(size(EEG_S1_85.epoch,2)) ' epochs']); 
disp([EEG_S2_85.filename ' has ' num2str(size(EEG_S2_85.epoch,2)) ' epochs']); 
disp([EEG_T_85.filename ' has ' num2str(size(EEG_T_85.epoch,2)) ' epochs']); 
end

